#include "lbvh.h"

#include <thrust\sequence.h>
#include <thrust\sort.h>
#include "build_kernels.h"
#include "update_kernels.h"

#include "collision_detection_kernels.h"
#include <base\cuda_device.cuh>

#include <physics\dynamics_kernel.cuh>
#include <collision\narrow_phase.h>

#include <iostream>
#include <time.h>
namespace gc {
	LBvh::LBvh(): primSize(-1), d_prims(nullptr) {
		// register build kernels
		MCudaDevice& CudaDevice = MCudaDevice::getInstance();
		CudaDevice.registerKernel(buildBVs);
		CudaDevice.registerKernel(calcSceneBV_1, cudaFuncCachePreferShared);
		CudaDevice.registerKernel(calcSceneBV_2, cudaFuncCachePreferShared);
		CudaDevice.registerKernel(calcMortonCodes_32b);
		CudaDevice.registerKernel(reorderLeaves);
		CudaDevice.registerKernel(calcBuildMetrics);
		CudaDevice.registerKernel(buildInteriors);
		CudaDevice.registerKernel(calcLeafEscape);
		CudaDevice.registerKernel(calcInteriorEscape);
		CudaDevice.registerKernel(calcInteriorNewIdx);
		CudaDevice.registerKernel(reorderInteriors);

		// register broad phase CD kernel
		CudaDevice.registerKernel(genBroadPhaseCP);

		// register update kernels
		CudaDevice.registerKernel(refitInteriors);

		// others
		CudaDevice.registerKernel(resetBVs);
	}

	LBvh::~LBvh() {
		this->free();
	}

	void LBvh::loadPrimitives(SimpleSphere *d_prims,int size) {
		this->primSize = size;
		this->d_prims = d_prims;

		leaves.malloc(primSize);
		interiors.malloc(primSize - 1);
		leaves_copy.malloc(primSize);
		interiors_copy.malloc(primSize - 1);
	}

	SimpleSphere* LBvh::getDevicePrimsPtr() {
		return d_prims;
	}

	void LBvh::free() {
		leaves.free();
		interiors.free();
		leaves_copy.free();
		interiors_copy.free();
	}

	void LBvh::build() {

		MCudaDevice& CudaDevice = MCudaDevice::getInstance();
		// calculate the BVs of all primitives
		CudaDevice.launchKernel(primSize, 0, buildBVs, primSize, leaves_copy, getConstPtr(d_prims));

		// using parallel reduce algorithm to calculate scene BV
		auto p  = CudaDevice.getOptLaunchParams(primSize, sizeof(BV), calcSceneBV_1);
		int gs = p.first, bs = p.second;
		thrust::device_vector<BV> d_bv_out(gs);
		// phase 1
		calcSceneBV_1 << <gs, bs, sizeof(BV)*bs >> > (primSize, leaves, getRawPtr(d_bv_out));
		getLastCudaError("calcSceneBV_1 kernel failure");
		
		// recursively do phase 2 until final result is computed (gridSize=1)
		while (gs != 1) {
			int size = gs;
			p = CudaDevice.getOptLaunchParams(size, sizeof(BV), calcSceneBV_2);
			gs = p.first, bs = p.second;
			calcSceneBV_2 << <gs, bs, sizeof(BV)*bs >> > (size, getRawPtr(d_bv_out), getRawPtr(d_bv_out));
			cudaDeviceSynchronize();
			getLastCudaError("calcSceneBV_2 kernel failure");
		}
		/*
		if (debug) {
			BV h_scene_bv;
			checkCudaErrors(GPU2CPU(&h_scene_bv, getRawPtr(d_bv_out), sizeof(BV)));
			std::cout << h_scene_bv.toString() << std::endl;
		}
		*/
		// calculate float32 morton code for each BV
		CudaDevice.launchKernel(primSize, 0, calcMortonCodes_32b, primSize, leaves_copy, getRawPtr(d_bv_out));
		
		// reorder leaves using morton code
		thrust::device_vector<int> mapidx(primSize);
		checkThrustErrors(thrust::sequence(mapidx.begin(), mapidx.end()));
		thrust::sort_by_key(getDevicePtr(leaves_copy.morton), getDevicePtr(leaves_copy.morton + primSize) , mapidx.begin());
		//thrust::host_vector<uint> mortons(primSize);
		//thrust::copy(getDevicePtr(leaves.morton), getDevicePtr(leaves.morton + primSize), mortons.begin());
		//for (int i = 0; i < primSize; ++i) printf("%u\n", mortons[i]);
		CudaDevice.launchKernel(primSize, 0, reorderLeaves, primSize, leaves_copy, leaves, getRawPtr(mapidx));
		cudaDeviceSynchronize();
		getLastCudaError("reorderLeaves kernel failure");

		checkCudaErrors(cudaMemset(leaves.segLen, 0, sizeof(int)*primSize));
		checkCudaErrors(cudaMemset(leaves.segOffset, 0, sizeof(int)*primSize));

		// calculate metrics
		thrust::device_vector<uint> metrics(primSize - 1);
		CudaDevice.launchKernel(primSize - 1, 0, calcBuildMetrics, primSize - 1, leaves, getRawPtr(metrics));

		// build interior nodes
		//CudaDevice.launchKernel(primSize - 1, 0, resetBVs, primSize - 1, interiors.bv);
		thrust::device_vector<int> visitCount(primSize - 1, 0);
		thrust::device_vector<int> leftLeafCount(primSize - 1, 0);
		InteriorBuildAid aid{ getRawPtr(metrics),getRawPtr(visitCount),getRawPtr(leftLeafCount) };
		CudaDevice.launchKernel(primSize, 0, buildInteriors, primSize, leaves, interiors_copy, aid);

		// calculate escape index for traversal
		CudaDevice.launchKernel(primSize, 0, calcLeafEscape, primSize, leaves, interiors_copy);
		CudaDevice.launchKernel(primSize - 1, 0, calcInteriorEscape, primSize - 1, leaves, interiors_copy);

		
		// reorder interiors
		/*
		thrust::host_vector<int> h_segLen(primSize);
		thrust::copy(getDevicePtr(leaves.segLen), getDevicePtr(leaves.segLen) + primSize, h_segLen.begin());
		for (int i = 0; i < primSize; ++i) {
			fprintf(stdout, "segLen[%d]: %d\n", i, h_segLen[i]);
		}
		*/
		thrust::exclusive_scan(getDevicePtr(leaves.segLen),getDevicePtr(leaves.segLen)+primSize,getDevicePtr(leaves.segOffset));
		/*
		thrust::host_vector<int> h_segoffset(primSize);
		thrust::copy(getDevicePtr(leaves.segOffset), getDevicePtr(leaves.segOffset) + primSize, h_segoffset.begin());
		for (int i = 0; i < primSize; ++i) {
			fprintf(stdout, "segOffset[%d]: %d\n", i,h_segoffset[i]);
		}
		*/
		CudaDevice.launchKernel(primSize - 1, 0, calcInteriorNewIdx, primSize - 1, leaves, interiors_copy, getRawPtr(leftLeafCount), getRawPtr(mapidx));
		CudaDevice.launchKernel(primSize-1, 0, reorderInteriors, primSize-1, leaves, interiors_copy, interiors, getRawPtr(mapidx));
		cudaDeviceSynchronize();
		getLastCudaError("reorderInteriors kernel failure");
		//fprintf(stdout, "reorderInteriors kernel success\n");
		
	}

	void LBvh::testBuild() {
		MCudaDevice& CudaDevice = MCudaDevice::getInstance();
		CudaDevice.launchKernel(primSize, 0, buildBVs, primSize, leaves, getConstPtr(d_prims));

		uint mortons[] = { 2,3,5,10,12,20,23,26 };
		checkCudaErrors(CPU2GPU(leaves.morton, mortons, primSize*sizeof(uint)));
		// calculate metrics
		thrust::device_vector<uint> metrics(primSize - 1);
		CudaDevice.launchKernel(primSize - 1, 0, calcBuildMetrics, primSize - 1, leaves, getRawPtr(metrics));

		// build interior nodes
		thrust::device_vector<int> visitCount(primSize - 1, 0);
		thrust::device_vector<int> leftLeafCount(primSize - 1, 0);
		InteriorBuildAid aid{ getRawPtr(metrics),getRawPtr(visitCount),getRawPtr(leftLeafCount) };
		CudaDevice.launchKernel(primSize, 0, buildInteriors, primSize, leaves, interiors, aid);

		// calculate escape index for traversal
		CudaDevice.launchKernel(primSize, 0, calcLeafEscape, primSize, leaves, interiors);
		CudaDevice.launchKernel(primSize - 1, 0, calcInteriorEscape, primSize-1, leaves, interiors);

		// reorder interiors
		thrust::exclusive_scan(getDevicePtr(leaves.segLen), getDevicePtr(leaves.segLen + primSize), getDevicePtr(leaves.segOffset));
		LBvhInteriors interiors_sorted;
		interiors_sorted.malloc(primSize - 1);
		thrust::device_vector<int> mapidx(primSize);
		CudaDevice.launchKernel(primSize - 1, 0, calcInteriorNewIdx, primSize - 1, leaves, interiors, getRawPtr(leftLeafCount), getRawPtr(mapidx));
		CudaDevice.launchKernel(primSize, 0, reorderInteriors, primSize, leaves, interiors, interiors_sorted, getRawPtr(mapidx));
		cudaDeviceSynchronize();
		getLastCudaError("reorderInteriors kernel failure");
		interiors.free();
		this->interiors = interiors_sorted;
	}

	void LBvh::update() {
		MCudaDevice& CudaDevice = MCudaDevice::getInstance();
		// calculate the BVs of all primitives
		CudaDevice.launchKernel(primSize, 0, buildBVs, primSize, leaves, getConstPtr(d_prims));
		//CudaDevice.launchKernel(primSize - 1, 0, resetBVs, primSize - 1, interiors.bv);
		//interiors.bv.reset();
		thrust::device_vector<int> visitCount(primSize-1, 0);
		CudaDevice.launchKernel(primSize, 0, refitInteriors, primSize, leaves, interiors, getRawPtr(visitCount));
	}

	void LBvh::broadPhaseCollison(CollisionPair *d_cp, int *size) {
		//thrust::device_vector<int> d_size(1, 0);
		MCudaDevice& CudaDevice = MCudaDevice::getInstance();
		int *d_size = nullptr;
		checkCudaErrors(GPUMalloc(d_size, sizeof(int)));
		checkCudaErrors(cudaMemset(d_size, 0, sizeof(int)));
		CudaDevice.launchKernel(primSize, 0, genBroadPhaseCP, primSize, leaves, interiors, d_cp, d_size);
		cudaDeviceSynchronize();
		getLastCudaError("genBroadPhaseCP kernel failure");
		checkCudaErrors(GPU2CPU(size, d_size, sizeof(int)));
		cudaFree(d_size);

	}
	
	BouncingBallSimulation::BouncingBallSimulation(SimpleBox &sceneBox, FILE* fLog)
		:n_prims(0), dt(0.01), frame(0), d_rigids(nullptr),
		kernelTime(0),cdTime(0),buildTime(0),narrowphaseTime(0), updateTime(0),
		fLog(fLog)
	{
		this->sceneBox = sceneBox;
		lx = sceneBox.halfSize.x * 2;
		ly = sceneBox.halfSize.y * 2;
		lz = sceneBox.halfSize.z * 2;
		base.x = sceneBox.center.x - sceneBox.halfSize.x;
		base.y = sceneBox.center.y - sceneBox.halfSize.y;
		base.z = sceneBox.center.z - sceneBox.halfSize.z;

		const int cpCapacity = 1e6;
		d_cp_total.resize(cpCapacity);

		MCudaDevice& CudaDevice = MCudaDevice::getInstance();
		CudaDevice.registerKernel(updateDynamics);
		CudaDevice.registerKernel(narrowPhase<SimpleSphere>);
	}

	BouncingBallSimulation::~BouncingBallSimulation() {
		//if (d_balls) checkCudaErrors(cudaFree(d_balls));
		if (d_rigids) checkCudaErrors(cudaFree(d_rigids));
		if (fLog != stderr && fLog != NULL) fclose(fLog);
	}

	void BouncingBallSimulation::setOutputPath(char *name) {
		fLog = fopen(name, "w");
	}

	void BouncingBallSimulation::setTimePerFrame(Float dt) { this->dt = dt; }

	void BouncingBallSimulation::generateBalls(int number,int maxVelocity) {
		if (fLog) fprintf(fLog, "scene: (%f,%f,%f) (%f,%f,%f)\n", base.x, base.y, base.z, base.x + lx, base.y + ly, base.z + lz);
		int sqrtn = int(sqrt(number));
		int nx = sqrtn + 3, nz = sqrtn + 3;
		n_prims = sqrtn * sqrtn;  // estimated size (upper bound)

		Float gridX = lx / nx;
		Float gridZ = lz / nz;

		Float radiusL = fmin(gridX, gridZ)*0.2;
		Float radiusR = radiusL * 2.5;

		//Float posyL = sceneBox.center.y;
		//Float posyR = sceneBox.center.y + sceneBox.halfSize.y - radiusR;
		Float posyL = radiusR;
		Float posyR = sceneBox.center.y;

		Float velocityL = 0, velocityR = maxVelocity;
		Float mass = 5;
		auto material = MPhys::PhysicalMaterial::Iron;

		// one in every grid
		h_balls.reserve(n_prims);
		//std::vector<SimpleSphere> balls(n_prims);
		std::vector<MPhys::RigidBody> rigids;
		rigids.reserve(n_prims);

		if (!debug) srand(time(0));
		else srand(1);
		for (Float x = base.x + gridX; x < base.x + lx - 2 * gridX; x += gridX) {
			for (Float z = base.z + gridZ; z < base.z + lz - 2 * gridZ; z += gridZ) {
				Float radius = genRand(radiusL, radiusR);
				Float posx = genRand(x, x + gridX);
				Float posy = genRand(posyL, posyR);
				Float posz = genRand(z, z + gridZ);
				Point3f pos{ posx,posy,posz };

				if (fLog) fprintf(fLog, "(%f,%f,%f) %f\n", posx, posy, posz, radius);

				Float speed = genRand(velocityL, velocityR);
				Point3f v = uniformSampleDirection()*speed;

				h_balls.push_back(SimpleSphere(pos, radius));
				auto rigid = MPhys::RigidBody(mass, material);
				rigid.v = v;
				rigid.position = pos;
				rigids.emplace_back(rigid);
			}
		}
		n_prims = h_balls.size();  // this is the real size
		if (fLog) fprintf(fLog, "nPrims: %d\n", n_prims);

		d_balls.resize(n_prims);
		//checkCudaErrors(GPUMalloc(d_balls, sizeof(SimpleSphere)*n_prims));
		//checkCudaErrors(CPU2GPU(getRawPtr(d_balls), h_balls.data(), sizeof(SimpleSphere)*n_prims));
		d_balls = h_balls;
		checkCudaErrors(GPUMalloc(d_rigids, sizeof(MPhys::RigidBody)*n_prims));
		checkCudaErrors(CPU2GPU(d_rigids, rigids.data(), sizeof(MPhys::RigidBody)*n_prims));

		lbvh.loadPrimitives(getRawPtr(d_balls), n_prims);
	}

	void BouncingBallSimulation::produceFrameBrutal() {
		if (fLog) fprintf(fLog, "frame %d\n", ++frame);
		MCudaDevice& CudaDevice = MCudaDevice::getInstance();
		Point3f boxMax{ base.x + lx,base.y + ly,base.z + lz };

		// tick
		clock_t start = clock();

		CudaDevice.launchKernel(n_prims, 0, updateDynamics, n_prims, getRawPtr(d_balls), d_rigids, dt, base, boxMax);
		BV_SOA bvs;
		bvs.malloc(n_prims);
		CudaDevice.launchKernel(n_prims, 0, buildBV_SOA, n_prims, bvs, getConstPtr(getRawPtr(d_balls)));

		int *d_CPsize = nullptr;
		checkCudaErrors(GPUMalloc(d_CPsize, sizeof(int)));
		checkCudaErrors(cudaMemset(d_CPsize, 0, sizeof(int)));
		clock_t cdStart = clock();
		CudaDevice.launchKernel(n_prims, 0, brutalForce, n_prims, bvs, getRawPtr(d_cp_total), d_CPsize);
		cudaDeviceSynchronize();
		cdTime += Float(clock() - cdStart);

		int h_cpSize;
		checkCudaErrors(GPU2CPU(&h_cpSize, d_CPsize, sizeof(int)));
		cudaFree(d_CPsize);

		if (frame % 100 == 0) {
			if (fLog) fprintf(stdout, "frame %d: BP collisions = %d\n", frame, h_cpSize);
		}

		clock_t npStart = clock();
		CudaDevice.launchKernel(h_cpSize, 0, narrowPhase<SimpleSphere>, h_cpSize, getRawPtr(d_cp_total), getRawPtr(d_balls), d_rigids);
		///< for output
		cudaDeviceSynchronize();

		// tock
		clock_t end = clock();
		narrowphaseTime += Float(end - npStart);
		kernelTime += Float(end - start);

		h_balls = d_balls;
		//checkCudaErrors(GPU2CPU(h_balls.data(), getRawPtr(d_balls), sizeof(SimpleSphere)*n_prims));
		for (auto &ball : h_balls) {
			if (fLog) fprintf(fLog, "%f %f %f\n", ball.center.x, ball.center.y, ball.center.z);
		}
	}

	void BouncingBallSimulation::produceFrame() {
		/*
		1. move every ball based on Newton physics
		2. pull back balls stuck in the wall
		3. build LBVH and find collision pairs
		4. do narrow phase interaction among CPs
		*/
		if (fLog) fprintf(fLog, "frame %d\n", ++frame);
		MCudaDevice& CudaDevice = MCudaDevice::getInstance();
		Point3f boxMax{ base.x + lx,base.y + ly,base.z + lz };

		// tick
		clock_t start = clock();

		CudaDevice.launchKernel(n_prims, 0, updateDynamics, n_prims, getRawPtr(d_balls), d_rigids, dt, base, boxMax);
		
		/*
		if (frame % 10 == 1) {
			clock_t buildStart = clock();
			lbvh.build();
			buildTime += Float(clock() - buildStart);
		}
		else { 
			clock_t updateStart = clock();
			lbvh.update(); 
			updateTime += Float(clock() - updateStart);
		}
		*/
		clock_t buildStart = clock();
		lbvh.build();
		buildTime += Float(clock() - buildStart);
		int h_cpSize;
		clock_t cdStart = clock();
		lbvh.broadPhaseCollison(getRawPtr(d_cp_total), &h_cpSize);
		cdTime += Float(clock() - cdStart);

		if (frame % 100 == 0) {
			if (fLog) fprintf(stdout, "frame %d: BP collisions = %d\n", frame, h_cpSize);
		}

		clock_t npStart = clock();
		CudaDevice.launchKernel(h_cpSize, 0, narrowPhase<SimpleSphere>, h_cpSize, getRawPtr(d_cp_total), getRawPtr(d_balls), d_rigids);
		///< for output
		cudaDeviceSynchronize();

		// tock
		clock_t end = clock();
		narrowphaseTime += Float(end - npStart);
		kernelTime += Float(end - start);

		h_balls = d_balls;
		//checkCudaErrors(GPU2CPU(h_balls.data(), getRawPtr(d_balls), sizeof(SimpleSphere)*n_prims));
		for (auto &ball : h_balls) {
			if (fLog) fprintf(fLog, "%f %f %f\n", ball.center.x, ball.center.y, ball.center.z);
		}
	}

	Float BouncingBallSimulation::timePerFrame() { return kernelTime / frame; }

	Float BouncingBallSimulation::timePerBroadPhase() { return cdTime / frame; }

	void BouncingBallSimulation::printTimeInfo() {
		fprintf(stdout, "Performance of %d-body build: %.3f ms/frame\n", n_prims, buildTime / (frame));
		fprintf(stdout, "Performance of %d-body udpate: %.3f ms/frame\n", n_prims, updateTime / (0.9*frame));
		fprintf(stdout, "Performance of %d-body broad phase: %.3f ms/frame\n", n_prims, cdTime / frame);
		fprintf(stdout, "Performance of %d-body narrow phase: %.3f ms/frame\n", n_prims, narrowphaseTime / frame);
		fprintf(stdout, "Performance of %d-body kernels: %.3f ms/frame\n", n_prims, kernelTime / frame);
	}
	
}